# [18]
import json

from dotenv import load_dotenv, find_dotenv
from openai import OpenAI

_ = load_dotenv(find_dotenv())

client = OpenAI()


# 只有 gpt-4 能跑动思维树。实验室不支持 gpt-4，自行实验请在本地运行

def get_completion(prompt, model="gpt-4", temperature=0):
    messages = [{"role": "user", "content": prompt}]
    response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=temperature  # 模型输出的随机性，0 表示随机性最小
    )
    return response.choices[0].message.content


def performance_analyser(text):
    prompt = (f"{text}\n请根据以上成绩，分析候选人在速度、耐力、力量三方面素质的分档。分档包括："
              f"\n以JSON格式输出，其中key为素质名，value为以数值表示的分档。")
    response = get_completion(prompt)
    return json.loads(response)


def possible_sports(talent, category):
    prompt = f"需要{talent}强的{category}运动有哪些。给出10个例子，以array形式输出。确保"
    response = get_completion(prompt, temperature=0.8)
    return json.loads(response)


def evaluate(sports, talent, value):
    prompt = (f"分析{sports}运动树{talent}方面素质的要求：强（3），中（2），弱（1）"
              f"\n直接输出档位数字。输出只包含数字。")
    response = get_completion(prompt)
    val = int(response)
    print(f"{sports}: {talent} {val} {value >= val}")
    return value >= val


def report_generator(name, performance, talents, sports):
    level = ['弱', '中', '强']
    _talents = {k: level[v - 1] for k, v in talents.items()}
    prompt = f"已知{name}{performance}\n身体素质：{_talents}。\n生成一篇{name}适合{sports}"
    response = get_completion(prompt, model="gpt-3.5-turbo")
    return response


name = "小明"
performance = "100 米跑成绩：10.5 秒，1500 米跑成绩：3 分 20 秒，铅球成绩：12 米。他适合参加哪些搏击运动训练"
category = "搏击"

talents = performance_analyser(name + performance)
print("===talents===")
print(talents)
